# -*- coding: utf-8 -*-

import os
import time
import cv2
from statistics import median
import numpy as np
import tensorflow as tf
from metro.vgg_train import vgg16_network


def read_labeled_points(label_file_path):
    with open(label_file_path) as label_file:
        str_ground_truths = label_file.readlines()
    labeled_points = [tuple([int(e) for e in line.strip().split(" ")]) for line in str_ground_truths]
    return labeled_points


def show_ground_truth(ground_truth, name="color_map"):
    min_value, max_value = np.min(ground_truth), np.max(ground_truth)
    img = ((ground_truth - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.resize(img, (img.shape[1]*4, img.shape[0]*4))
    color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    cv2.imshow(name, color_map)


def generate_mask():
    image_mask = np.zeros(shape=(520, 960, 3), dtype=np.uint8)
    image_mask[...] = 255
    cv2.line(image_mask, (325, 0), (0, 475), (0, 0, 0), 2)
    cv2.line(image_mask, (500, 0), (960, 250), (0, 0, 0), 2)
    # cv2.line(image_mask, (500, 0), (960, 350), (0, 0, 0), 2)
    cv2.floodFill(image_mask, None, (160, 230), (0, 0, 0))
    cv2.floodFill(image_mask, None, (850, 130), (0, 0, 0))
    return image_mask


def read_image_and_ground_truth():
    folder = "/home/lijun/Dataset/metro/test/2016_09_29_17_19"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    label_files_name = [name.replace("jpg", "txt") for name in images_name]
    image_mask = generate_mask()
    for image_name, label_file_name in zip(images_name, label_files_name):
        image_path = os.path.join(folder, image_name)
        image = cv2.imread(image_path)
        image = image[40:, :, :]
        image = cv2.resize(image, dsize=(image.shape[1]//2, image.shape[0]//2))
        image &= image_mask
        # cv2.line(image, (0, image.shape[0]//5), (image.shape[1], image.shape[0]//5), (0, 255, 255), 2)
        cv2.imshow("image", image)
        image = image.astype(np.float32, copy=False)
        image /= 255
        image = np.expand_dims(image, axis=0)
        labeled_points = read_labeled_points(os.path.join(folder, label_file_name))
        yield image, len(labeled_points), image_name


def main():
    predict_counts, ground_truth_counts = [], []
    images_name = []
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32, shape=[None, 520, 960, 3])
        inference_op = vgg16_network(input_x)
        saver = tf.train.Saver()
        with tf.Session() as session:
            saver.restore(session, "../model/metro/vgg-25140")
            for image, ground_truth, image_name in read_image_and_ground_truth():
                inference = session.run(inference_op, feed_dict={input_x: image})
                inference = inference[0][:, :, 0]
                top_half = np.sum(inference[0:inference.shape[0]//4, :]) / 1000
                medium_half = np.sum(inference[inference.shape[0]//4: inference.shape[0]//4*3, :]) / 1000
                bottom_half = np.sum(inference[inference.shape[0]//4*3:, :]) / 1000
                top_half *= 0.9
                bottom_half *= 1.2
                predict_count = int(top_half + medium_half + bottom_half)
                predict_counts.append(predict_count)
                ground_truth_counts.append(ground_truth)
                images_name.append(image_name)
            with open("./test_result/2016_09_29_17_19.txt", "w") as result_file:
                for name, predict, ground_truth in zip(images_name, predict_counts, ground_truth_counts):
                    str_line = name + " " + str(predict) + " " + str(ground_truth) + "\n"
                    result_file.writelines(str_line)


if __name__ == "__main__":
    main()
